# File: run_agent.py (Should be correct for SSE)

import asyncio
import sys
from pathlib import Path
from dotenv import load_dotenv
import logging
import os
import traceback # 导入 traceback
from typing import Optional

# --- 路径和环境设置 --- (保持不变)
project_root = Path(__file__).parent
src_path = project_root / "src"
if str(src_path) not in sys.path:
    sys.path.insert(0, str(src_path))
    print(f"[run_agent.py] Added src path to sys.path: {src_path}")

env_path = project_root / ".env"
if env_path.exists():
    loaded = load_dotenv(dotenv_path=env_path)
    print(f"[run_agent.py] Loaded environment variables from {env_path}: {loaded}")
else:
    print(f"[run_agent.py] .env file not found at {env_path}, using system environment.")

# --- 日志设置 --- (保持不变)
logger: Optional[logging.Logger] = None
try:
    from src.core.logging_setup import setup_logging_from_config, get_logger
    setup_logging_from_config()
    logger = get_logger("agent_runner")
    logger.info("Logging configured for agent runner.")
except Exception as e:
    logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(name)s - %(message)s")
    logger = logging.getLogger("agent_runner")
    logger.error(f"Failed to setup logging from config: {e}", exc_info=True)
    print(f"WARNING: Failed to setup logging from config, using basic logging: {e}", file=sys.stderr)
if logger is None:
    logger = logging.getLogger("agent_runner_fallback")


# --- 导入智能体和客户端组件 --- (保持不变)
try:
    from src.agent.orchestrator import run_task_loop
    from src.mcp_client import initialize_mcp_client, close_mcp_client
    from src.core.exceptions import MCPClientError, ConfigError, AgentError
except ImportError as e:
     logger.critical(f"Failed to import necessary agent or client modules: {e}", exc_info=True)
     print(f"ERROR: Failed to import necessary modules: {e}", file=sys.stderr)
     sys.exit(1)


async def agent_main_logic():
    """包含运行智能体任务循环的逻辑"""
    # *** 添加日志 ***
    assert logger is not None, "Logger should have been initialized before running agent logic"

    logger.info(">>> Entering agent_main_logic()...")
    try:
        initial_task_description = "访问百度，搜索'自动驾驶汽车'，然后总结第一个结果页面的主要内容"
        logger.info(f"Agent Task: {initial_task_description}")
        # 调用核心任务循环 (位于 src/agent/orchestrator.py)
        # *** 添加日志 ***
        logger.info(">>> Calling run_task_loop()...")
        await run_task_loop(initial_task_description)
        logger.info("<<< run_task_loop() finished.")
        logger.info("Agent main logic finished successfully.")
    except AgentError as e:
        # *** 使用 exc_info=True 记录完整 traceback ***
        logger.error(f"Agent execution failed within agent_main_logic: {e}", exc_info=True)
    except Exception as e:
        # *** 使用 exc_info=True 记录完整 traceback ***
        logger.critical(f"An unexpected error occurred during agent execution within agent_main_logic: {e}", exc_info=True)
    # *** 添加日志 ***
    logger.info("<<< Exiting agent_main_logic().")


async def run_with_sse(): # 函数名已修改
    """初始化 MCP 客户端 (SSE/TCP) 并运行智能体逻辑"""
    assert logger is not None, "Logger should have been initialized before running agent logic"

    logger.info("Attempting to run agent with SSE connection...")
    conn_type_check = os.getenv("MCP_CONNECTION_TYPE", "sse").lower()
    if conn_type_check != 'sse':
         logger.warning(f"MCP_CONNECTION_TYPE is '{conn_type_check}', but proceeding assuming SSE based on session_manager config.")

    mcp_session_initialized = False
    try:
        logger.info("Connecting to running MCP server via SSE...")
        # *** 添加日志 ***
        logger.info(">>> Calling initialize_mcp_client()...")
        await initialize_mcp_client()
        mcp_session_initialized = True
        # *** 添加日志 ***
        logger.info("<<< initialize_mcp_client() successful.")
        logger.info("MCP client connection initialized successfully.")

        # 运行主要的智能体逻辑
        # *** 添加日志 ***
        logger.info(">>> Calling agent_main_logic()...")
        await agent_main_logic()
        # *** 添加日志 ***
        logger.info("<<< agent_main_logic() returned.")

    except (MCPClientError, ConfigError) as e:
        # *** 修改日志级别和 exc_info ***
        logger.critical(f"Failed to initialize or maintain MCP connection (SSE): {e}", exc_info=False) # Keep critical, no full tb needed here usually
    except Exception as e:
        # *** 记录其他意外错误，确保记录 traceback ***
        logger.critical(f"An unexpected error occurred in run_with_sse (outside agent_main_logic?): {e}", exc_info=True) # Log full traceback
    finally:
        logger.info("Ensuring MCP client cleanup (SSE)...")
        await close_mcp_client()
        logger.info("MCP client cleanup finished (SSE).")


# --- 主执行入口 --- (保持不变)
if __name__ == "__main__":
    print("-" * 60)
    print("[run_agent.py] Starting agent execution using SSE server connection...")
    print("!!! Ensure the server is already running via 'python run.py' !!!")
    print("-" * 60)

    if logger is None:
        print("FATAL: Logger was not initialized. Exiting.", file=sys.stderr)
        sys.exit(1)

    try:
        asyncio.run(run_with_sse())
    except KeyboardInterrupt:
        print("\n[run_agent.py] KeyboardInterrupt received. Shutting down...")
        logger.info("KeyboardInterrupt received, initiating shutdown.")
        try:
            pass
        except Exception as cleanup_err:
             logger.error(f"Error during manual cleanup after KeyboardInterrupt: {cleanup_err}", exc_info=True)
    except Exception as main_err:
        print(f"\n[run_agent.py] CRITICAL ERROR in main execution: {main_err}", file=sys.stderr)
        traceback.print_exc(file=sys.stderr)
        # *** 确保记录 traceback ***
        logger.critical(f"CRITICAL ERROR in main agent execution: {main_err}", exc_info=True)

    print("-" * 60)
    print("[run_agent.py] Agent execution finished.")
    logger.info("Agent execution script finished.")
    print("-" * 60)